package com.ge.executor;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.lang.UUID;
import com.ge.executor.enums.LanguageCmdEnum;
import com.ge.executor.model.ExecuteMessage;
import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.async.ResultCallback;
import com.github.dockerjava.api.command.CreateContainerCmd;
import com.github.dockerjava.api.command.CreateContainerResponse;
import com.github.dockerjava.api.command.ExecCreateCmdResponse;
import com.github.dockerjava.api.model.Frame;
import com.github.dockerjava.api.model.HostConfig;
import com.github.dockerjava.api.model.StreamType;
import com.github.dockerjava.core.DockerClientBuilder;
import lombok.extern.slf4j.Slf4j;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;

/**
 * docker 沙箱
 */
@Slf4j
public class DockerSandbox {

    private static final DockerClient DOCKER_CLIENT = DockerClientBuilder.getInstance().build();

    /**
     * 执行代码
     *
     * @param languageCmdEnum 编程语言枚举
     * @param code            代码
     * @return {@link ExecuteMessage}
     */
    public static ExecuteMessage execute(LanguageCmdEnum languageCmdEnum, String code) {

        // 写入文件
        String userDir = System.getProperty("user.dir");
        String language = languageCmdEnum.getLanguage();
        String globalCodePathName = userDir + File.separator + "tempCode" + File.separator + language;
        // 判断全局代码目录是否存在，没有则新建
        File globalCodePath = new File(globalCodePathName);
        if (!globalCodePath.exists()) {
            boolean mkdir = globalCodePath.mkdirs();
            if (!mkdir) {
                log.info("创建全局代码目录失败");
            }
        }

        // 把用户的代码隔离存放
        String userCodeParentPath = globalCodePathName + File.separator + UUID.randomUUID();
        String userCodePath = userCodeParentPath + File.separator + languageCmdEnum.getSaveFileName();
        FileUtil.writeString(code, userCodePath, StandardCharsets.UTF_8);
        String containerId = createContainer(userCodePath);

        ExecuteMessage executeMessage;

        // 编译代码
        String[] compileCmd = languageCmdEnum.getCompileCmd();
        if (compileCmd != null) {
            executeMessage = execCmd(containerId, compileCmd);
            log.info("编译完成...");
            if (!executeMessage.isSuccess()) {
                // 清除文件
                cleanFileAndContainer(userCodeParentPath, containerId);
                return executeMessage;
            }
        }
        executeMessage = execCmd(containerId, languageCmdEnum.getRunCmd());
        log.info("运行完成...");

        // 清除文件
        cleanFileAndContainer(userCodeParentPath, containerId);
        return executeMessage;
    }

    /**
     * 创建容器
     *
     * @param codeFile
     * @return
     */
    private static String createContainer(String codeFile) {
        CreateContainerCmd containerCmd = DOCKER_CLIENT.createContainerCmd("codesandbox:latest");
        HostConfig hostConfig = new HostConfig();
        hostConfig.withMemory(60 * 1000 * 1000L);
        hostConfig.withMemorySwap(0L);
        hostConfig.withCpuCount(1L);

        CreateContainerResponse createContainerResponse = containerCmd
                .withHostConfig(hostConfig)
                .withNetworkDisabled(true)
                .withAttachStdin(true)
                .withAttachStdout(true)
                .withAttachStderr(true)
                .withTty(true)
                .exec();
        // 启动容器
        String containerId = createContainerResponse.getId();
        DOCKER_CLIENT.startContainerCmd(containerId).exec();

        // 将代码复制到容器中
        DOCKER_CLIENT.copyArchiveToContainerCmd(containerId)
                .withHostResource(codeFile)
                .withRemotePath("/box")
                .exec();
        return containerId;
    }

    /**
     * 执行命令
     *
     * @param containerId 容器 ID
     * @param cmd         CMD
     * @return {@link ExecuteMessage}
     */
    private static ExecuteMessage execCmd(String containerId, String[] cmd) {
        // 正常返回信息
        ByteArrayOutputStream resultStream = new ByteArrayOutputStream();
        // 错误信息
        ByteArrayOutputStream errorResultStream = new ByteArrayOutputStream();
        // 结果
        final boolean[] result = {true};

        try {
            // 处理容器执行结果
            ResultCallback.Adapter<Frame> frameAdapter = new ResultCallback.Adapter<Frame>() {
                @Override
                public void onNext(Frame frame) {
                    StreamType streamType = frame.getStreamType();
                    byte[] payload = frame.getPayload();
                    if (StreamType.STDERR.equals(streamType)) {
                        // 执行错误
                        result[0] = false;
                        try {
                            errorResultStream.write(payload);
                        } catch (IOException e) {
                            throw new RuntimeException(e);
                        }
                    } else {
                        result[0] = true;
                        try {
                            resultStream.write(payload);
                        } catch (IOException e) {
                            throw new RuntimeException(e);
                        }
                    }
                    super.onNext(frame);
                }
            };

            ExecCreateCmdResponse execCompileCmdResponse = DOCKER_CLIENT.execCreateCmd(containerId)
                    .withCmd(cmd)
                    .withAttachStderr(true)
                    .withAttachStdin(true)
                    .withAttachStdout(true)
                    .exec();
            String execId = execCompileCmdResponse.getId();
            DOCKER_CLIENT.execStartCmd(execId).exec(frameAdapter).awaitCompletion();
            return ExecuteMessage.builder()
                    .success(result[0])
                    .message(resultStream.toString())
                    .errorMessage(errorResultStream.toString())
                    .build();

        } catch (InterruptedException e) {
            log.info(e.getMessage());
            return ExecuteMessage.builder()
                    .success(false)
                    .errorMessage(e.getMessage())
                    .build();
        }
    }

    /**
     * 清理文件和容器
     *
     * @param userCodePath 用户代码路径
     * @param containerId  容器 ID
     */
    private static void cleanFileAndContainer(String userCodePath, String containerId) {
        // 清理临时目录
        FileUtil.del(userCodePath);

        DOCKER_CLIENT.stopContainerCmd(containerId).exec();
        DOCKER_CLIENT.removeContainerCmd(containerId).exec();
    }
}
